#include "detector/number_classifier.hpp"

#include <algorithm>
#include <iostream>

#include <fmt/color.h>
#include <fmt/core.h>

namespace detector {
NumberClassifier::NumberClassifier(const std::string &model_path)
    : net_ok_(false) {
  const std::string& path{model_path};
  try {
    net_ = cv::dnn::readNetFromONNX(path);
    net_.setPreferableBackend(cv::dnn::DNN_BACKEND_CUDA);
    net_.setPreferableTarget(cv::dnn::DNN_TARGET_CUDA_FP16);
    net_ok_ = true;
  } catch (const std::exception &e) {
    fmt::print(
        fg(fmt::color::red) | fmt::emphasis::bold | fmt::emphasis::italic,
        "\n\n\n读取模型{}失败\n"
        "    请尝试修改configures/default/params.yaml中的model_path\n\n\n",
        model_path);
    throw(e);
  }
}

std::tuple<int, float> NumberClassifier::classify(cv::Mat &roi) {
  // get input
  cv::Mat tmp;
  roi.convertTo(tmp, CV_32FC1);
  tmp = tmp / 255.0;
  // set input
  cv::Mat blob = cv::dnn::blobFromImage(tmp, 1., cv::Size(28, 28));
  net_.setInput(blob);
  // get output
  cv::Mat outputs = net_.forward();
  // process output
  float obj_pb = *outputs.begin<float>();
  if (fabs(obj_pb) > 0.75) {
    cv::Mat cls_score = cv::Mat(9, 1, CV_32FC1);
    std::copy(outputs.begin<float>() + 1, outputs.end<float>(),
              cls_score.begin<float>());
    float max_prob =
        *std::max_element(cls_score.begin<float>(), cls_score.end<float>());
    // do softmax
    cv::Mat softmax_prob;
    cv::exp(cls_score - max_prob, softmax_prob);
    float sum = static_cast<float>(cv::sum(softmax_prob)[0]);
    softmax_prob /= sum;
    // get highest confidence
    float confidence = 0;
    int id = 0;
    int idx = 0;
    for (auto iter = softmax_prob.begin<float>();
         iter != softmax_prob.end<float>(); ++iter) {
      if (*iter > confidence) {
        confidence = *iter;
        id = idx;
      }
      ++idx;
    }
    return std::make_tuple(id, confidence);
  } else {
    return std::make_tuple(-1, 1.0);
  }
}

} // namespace detector